"""
FastAPI WebSocket 服务端实现
提供WebSocket连接管理和消息广播功能
"""

from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles  # 用于托管静态文件
from fastapi.responses import HTMLResponse  # 返回HTML响应
from pydantic import BaseModel
import uvicorn  # ASGI服务器
from typing import Any, Dict, List  # 类型注解
import json  # JSON处理
import time  # 时间处理

# 创建FastAPI应用实例
app = FastAPI()


# 定义可接受的JSON消息结构
class ClientMessage(BaseModel):
    type: str  # 消息类型
    data: Dict[str, Any]  # 消息内容
    timestamp: int = None  # 可选时间戳
# WebSocket连接管理器
class ConnectionManager:
    def __init__(self):
        # 存储所有活跃的WebSocket连接
        self.active_connections: List[WebSocket] = []

    # 建立新连接
    async def connect(self, websocket: WebSocket):
        # 等待客户端连接
        await websocket.accept()
        # 将新连接添加到活跃连接列表
        self.active_connections.append(websocket)
        print(f"新客户端连接，当前连接数: {len(self.active_connections)}")

    # 断开连接
    def disconnect(self, websocket: WebSocket):
        # 从活跃连接列表中移除
        self.active_connections.remove(websocket)
        print(f"客户端断开，剩余连接数: {len(self.active_connections)}")

    # 向单个客户端发送消息
    async def send_personal_message(self, message: str, websocket: WebSocket):
        # 发送文本消息
        await websocket.send_text(message)

    # 向所有客户端广播消息
    async def broadcast(self, message: str):
        # 遍历所有活跃连接并发送消息
        for connection in self.active_connections:
            await connection.send_text(message)

# 创建连接管理器实例
manager = ConnectionManager()

# 测试用的HTML页面
@app.get("/")
async def get():
    return HTMLResponse("""
    <html>
        <head>
            <title>WebSocket测试</title>
        </head>
        <body>
            <h1>WebSocket测试服务端</h1>
            <p>请使用前端页面连接WebSocket</p>
        </body>
    </html>
    """)

# WebSocket端点
@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    # 建立连接
    await manager.connect(websocket)
    
    try:
        while True:
            # 接收客户端消息
            json_data = await websocket.receive_text()
            data =  json.loads(json_data)
            print(data)
            # 构造返回消息
            message = {
                "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
                "client": data["data"]["name"],
                "content": data["data"]["content"]
            }
            
            # 广播消息给所有客户端
            await manager.broadcast(json.dumps(message, ensure_ascii=False))
            
    except WebSocketDisconnect:
        # 处理客户端断开连接
        manager.disconnect(websocket)
        # 通知其他客户端有用户离开
        message = {
            "time": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()),
            "client": "系统",
            "content": "有用户离开了聊天室"
        }
        await manager.broadcast(json.dumps(message, ensure_ascii=False))

# 启动FastAPI应用
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)